
import numpy as np

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.transforms import Affine2D
from matplotlib.ticker import MultipleLocator, FormatStrFormatter

import dicom
from dicom.filereader import ReadFile
import sys

from show_blit_ import CrossHair 

# doseDICOMpare.py
"""


"""
class dicomDOSE:
    """ class to store DICOM dose information, as well as plot dose distribution"""
    def __init__(self,identifier):#initialise class members
        self.filename = "" 
        self.offset = 0
        self.doseplane = "XZ"
        self.identifier = identifier
        #
        #limits of arrays, ie., indicies
        self.xmin = 0
        self.xmax = 1000
        self.ymin = 0
        self.ymax = 1000
        self.zmin = 0
        self.zmax = 1000
        #limits of dimensions in mm
        self.Xmin = 0
        self.Xmin = 0
        self.Ymin = 0
        self.Ymax = 0
        self.Zmin = 0
        self.Zmax = 0
        #
        self.title = "reference"
        if(self.identifier==2) : self.title = "non-reference"
        self.cb = False
        self.ax = False
        self.profax = False #profile axes
        self.xvalues = None
        self.yvalues = None
        self.zvalues = None
        self.cross_hair = None
        self.crossHairsFrozen = False
        self.cross_hairX = None
        self.cross_hairY = None#coords of cross hairs can be saved
        self.profileDirection = "X"#which profile will be plot
        self.sliceCoord = None
        self.dataset = None
        self.isocentre = [0,0,0]#set the default isocentre
        
    def TwoDSlice(self):
		if(self.doseplane=="XZ"):
			offsetInRange = self.offset*self.ymax
			self.slice=self.dataset.pixel_array[:,offsetInRange,:]#ZYX, slice will be ZX
			self.sliceCoord = self.Ymax - self.offset*(self.Ymax-self.Ymin)#Y is out of screen
		elif(self.doseplane=="YZ"):
			offsetInRange = self.offset *self.xmax
			self.slice=self.dataset.pixel_array[:,:,offsetInRange]#ZYX, slice will be ZY
			self.sliceCoord = self.offset*(self.Xmax-self.Xmin)+self.Xmin
		else:
			offsetInRange = self.offset *self.zmax
			self.slice=self.dataset.pixel_array[offsetInRange,:,:]#z slice is 0th element, ZYX, slice will be YX
			self.sliceCoord = self.Zmax - self.offset*(self.Zmax-self.Zmin)#Y is out of screen
		#do the dose scaling
		self.slice=np.multiply(self.slice,self.dataset.DoseGridScaling)

        
#loking at x profiles for now
    def DrawSecondaryProfile(self,fig,xdata,ydata,primaryProfax):

        if(self.crossHairsFrozen) :
            xdata = self.cross_hairX
            ydata = self.cross_hairY

        if(self.doseplane=="XZ") :
            if(self.profileDirection=="X"):
                index = (self.Zmax-ydata) / (self.Zmax - self.Zmin) * self.zmax#value minus min value divided by
                primaryProfax.plot(self.xvalues,self.slice[index,:])
            else :
                index = (xdata - self.Xmin) / (self.Xmax - self.Xmin) * self.xmax#value minus min value divided by
                temparray = self.zvalues.tolist()
                primaryProfax.plot(temparray[::-1],self.slice[:,index])
        elif(self.doseplane=="YZ") :
            if(self.profileDirection=="X"):
                index = (self.Zmax-ydata) / (self.Zmax - self.Zmin) * self.zmax
                temparray = self.yvalues.tolist()
                primaryProfax.plot(temparray[::-1],self.slice[index,:])
            else :
                index = (self.Ymax - xdata) / (self.Ymax - self.Ymin) * self.ymax
                temparray = self.zvalues.tolist()
                primaryProfax.plot(temparray[::-1],self.slice[:,index])
        else : #XY plot
            if(self.profileDirection=="X"):
                index = (self.Ymax-ydata) / (self.Ymax - self.Ymin) * self.ymax
                primaryProfax.plot(self.xvalues,self.slice[index,:])
            else :
                index = (xdata - self.Xmin) / (self.Xmax - self.Xmin) * self.xmax
                temparray = self.yvalues.tolist()
                primaryProfax.plot(temparray[::-1],self.slice[:,index])

        fig.canvas.draw()

    def DrawProfile(self, fig, xdata, ydata):
    #if this is the primary profile, primaryProfax will be None
        #print 'ydata', ydata

        if(self.crossHairsFrozen) :
            xdata = self.cross_hairX
            ydata = self.cross_hairY

        if(self.profax) : self.profax.clear()        
        majorFormatter = FormatStrFormatter('%2.1f')

        #only add subplot if it's the primary dose distribution
        self.profax = fig.add_subplot("111", title="profile")
        self.profax.set_autoscaley_on(False)
        
        #based on ydata , determine index, for, XZ and YZ, this will be based on Zmin and Zmax limits
        if(self.doseplane=="XZ") : 
            if(self.profileDirection=="X"):
                index = (self.Zmax-ydata) / (self.Zmax - self.Zmin) * self.zmax#value minus min value divided by
                self.profax.plot(self.xvalues,self.slice[index,:]) 
                if (index > self.xmax) : index = self.xmax -1
                self.profax.axis([self.Xmin, self.Xmax, 0, self.doseMax*1.1])
                self.profax.set_xlabel('X, mm')
            else :
                index = (xdata - self.Xmin) / (self.Xmax - self.Xmin) * self.xmax#value minus min value divided by
                temparray = self.zvalues.tolist()
                self.profax.plot(temparray[::-1],self.slice[:,index])
                self.profax.axis([self.Zmin, self.Zmax, 0, self.doseMax*1.1])
                self.profax.set_xlabel('Z, mm')
        elif(self.doseplane=="YZ") : 
            if(self.profileDirection=="X"):
                index = (self.Zmax-ydata) / (self.Zmax - self.Zmin) * self.zmax
                temparray = self.yvalues.tolist()
                #wow, shit python syntaxxx
                self.profax.plot(temparray[::-1],self.slice[index,:])
                self.profax.axis([self.Ymax, self.Ymin, 0, self.doseMax*1.1])
                self.profax.set_xlabel('Y, mm')
            else :
                index = (self.Ymax - xdata) / (self.Ymax - self.Ymin) * self.ymax
                temparray = self.zvalues.tolist()
                self.profax.plot(temparray[::-1],self.slice[:,index])
                self.profax.axis([self.Zmax, self.Zmin, 0, self.doseMax*1.1])
                self.profax.set_xlabel('Z, mm')
        else : #XY plot
            if(self.profileDirection=="X"):
                index = (self.Ymax-ydata) / (self.Ymax - self.Ymin) * self.ymax
                self.profax.plot(self.xvalues,self.slice[index,:])
                self.profax.axis([self.Xmin, self.Xmax, 0, self.doseMax*1.1])
                self.profax.set_xlabel('X, mm')
            else :
                index = (xdata - self.Xmin) / (self.Xmax - self.Xmin) * self.xmax
                temparray = self.yvalues.tolist()
                self.profax.plot(temparray[::-1],self.slice[:,index])
                self.profax.axis([self.Ymin, self.Ymax, 0, self.doseMax*1.1])
                self.profax.set_xlabel('Y, mm')
            
        self.profax.yaxis.set_major_formatter(majorFormatter)
        self.profax.yaxis.set_major_locator(MaxNLocator(5))
        self.profax.xaxis.set_major_locator(MaxNLocator(10))

        self.profax.set_ylabel('Dose, cGy')
        fig.canvas.draw()
        #

    def DrawFigure(self, fig):
        if(self.ax) : self.ax.clear()
        self.ax = fig.add_subplot("111", title=self.title, aspect='equal')
        #now show it and modify limits / axis labels depending on mode
        #visualise in ROOM coord system
        if(self.doseplane=="XZ") : 
            cax = self.ax.imshow(np.flipud(self.slice), extent=(self.Xmin, self.Xmax, self.Zmax, self.Zmin),vmin=0,vmax=self.doseMax, aspect='auto')
            self.ax.set_xlabel('X, mm')
            self.ax.set_ylabel('Z, mm')
        elif(self.doseplane=="YZ") : 
            cax = self.ax.imshow(np.fliplr(self.slice), extent=(self.Ymin, self.Ymax, self.Zmin, self.Zmax),vmin=0,vmax=self.doseMax, aspect='auto')
            self.ax.set_xlabel('Y, mm')
            self.ax.set_ylabel('Z, mm')
        else : #XY
            cax = self.ax.imshow(self.slice, extent=(self.Xmin, self.Xmax, self.Ymin, self.Ymax),vmin=0,vmax=self.doseMax, aspect='auto')#nb. y direction flipped
            self.ax.set_xlabel('X, mm')
            self.ax.set_ylabel('Y, mm')
        if (not(self.cb)):
            self.cb = fig.colorbar(cax,format='%3.2f')
            self.cb.set_label('Dose, cGy')
        fig.canvas.draw()
        #if(self.crossHairsFrozen) : self.RedrawCrossHairs()
        self.ax_background = fig.canvas.copy_from_bbox(self.ax.bbox)
        if self.cross_hair is None :
            self.cross_hair = CrossHair(self.ax, horizontal=True, vertical=True)
        #fig.canvas.draw()

    def update(self, fig, event):
        x = None
        y = None
        if(event):
            x = event.xdata
            y = event.ydata
        #if we have frozen cross hairs
        if(self.crossHairsFrozen) :
            x = self.cross_hairX
            y = self.cross_hairY
        if x is None or y is None: return        
        fig.canvas.restore_region(self.ax_background)        
        self.cross_hair.redraw(x=x, y=y)
        fig.canvas.blit(self.ax.bbox)

    def HideCrossHairs(self,fig):
        self.cross_hair.hide()
        fig.canvas.draw()

    def ShowCrossHairs(self,fig):
        self.cross_hair.show()

    def SaveCrossHairs(self,x,y):
        self.cross_hairX = x
        self.cross_hairY = y

    def RedrawCrossHairs(self):
        self.cross_hair.redraw(x=self.cross_hairX, y=self.cross_hairY)


#method to associate plan iwth dose distribution, needde to shift isocentre
    def ReadDICOMPlan(self,filename):
        self.plan = ReadFile(filename)
        self.isocentre = self.plan.Beams[0].ControlPoints[0].IsocenterPosition
        print "isocentre ", self.isocentre
        self.calculateLimits()#need to readjust limits

    def ReadDICOM(self, filename):
        self.filename = filename
        self.dataset = ReadFile(filename)
        #set limits of dose distribution
        self.zmax = self.dataset.pixel_array.shape[0]#this is actually z
        self.ymax = self.dataset.pixel_array.shape[1]#this is rows, ie,. y
        self.xmax = self.dataset.pixel_array.shape[2]#this is  coloumns, ie,.x
        #scale data by scaling factor
        print "scaling factor ", self.dataset.DoseGridScaling
        self.calculateLimits()

    def calculateLimits(self):
        self.doseMax = self.dataset.pixel_array.max()
        self.doseMax = self.doseMax*self.dataset.DoseGridScaling
        print 'max ', self.doseMax
        print self.dataset.ImagePositionPatient

        self.Xmin = self.dataset.ImagePositionPatient[0] - 5
        self.Ymin = self.dataset.ImagePositionPatient[1] - 5
        self.Zmin = self.dataset.ImagePositionPatient[2]
        print self.dataset.PixelSpacing

        print self.xmax*self.dataset.PixelSpacing[0]
        print self.ymax*self.dataset.PixelSpacing[1]

        self.Xmax = self.xmax*self.dataset.PixelSpacing[0] + self.Xmin
        self.Ymax = self.ymax*self.dataset.PixelSpacing[1] + self.Ymin
        #print self.dataset.GridFrameOffsetVector

        zSlices = self.dataset.NumberofFrames
        #print zSlices
        self.Zmax = self.dataset.GridFrameOffsetVector[self.dataset.NumberofFrames-1] + self.Zmin 
        print self.Xmin, self.Xmax, self.Ymin, self. Ymax, self.Zmin, self.Zmax
        print self.xmin, self.xmax, self.ymin, self. ymax, self.zmin, self.zmax
        self.xvalues = np.linspace(self.Xmin,self.Xmax,self.xmax)
        self.yvalues = np.linspace(self.Ymin,self.Ymax,self.ymax)
        self.zvalues = np.linspace(self.Zmin,self.Zmax,self.zmax)
        self.xvalues = self.xvalues - self.isocentre[0]
        self.yvalues = self.yvalues - self.isocentre[1]
        self.zvalues = self.zvalues - self.isocentre[2]
        #print self.xvalues
        #print self.yvalues
        #print self.zvalues
        #adjust limits based on isocentre position
        self.Xmin = self.Xmin - self.isocentre[0]
        self.Ymin = self.Ymin - self.isocentre[1]
        self.Zmin = self.Zmin - self.isocentre[2]
        self.Xmax = self.Xmax - self.isocentre[0]
        self.Ymax = self.Ymax - self.isocentre[1]
        self.Zmax = self.Zmax - self.isocentre[2]



